#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import sys
import glob
import tarfile
import filecmp
import tempfile
import subprocess
from bbki._util import Util
from bbki._util import TempChdir
from bbki._util import PhysicalDiskMounts


def usage(onErr=False):
    f = sys.stderr if onErr else sys.stdout
    exitCode = 1 if onErr else 0
    print("Usage: %s [initramfs-file] add-bin <file>" % (sys.argv[0]), file=f)
    print("       %s [initramfs-file] add-file <file> <target>" % (sys.argv[0]), file=f)
    sys.exit(exitCode)


def die(msg):
    print(msg, file=sys.stderr)
    sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        usage(True)

    if sys.argv[1] in ["-h", "--help"]:
        usage()

    initramfsFullfn = None
    initramfsFilesFullfn = None
    op = None
    if True:
        i = 1
        if sys.argv[i] not in ["add-bin", "add-file"]:
            initramfsFullfn = sys.argv[1]
            if not initramfsFullfn.startswith("/boot/initramfs-"):
                die("Invalid initramfs file.")
            if not os.path.exists(initramfsFullfn):
                die('File "%s" does not exist.' % (initramfsFullfn))
            postfix = initramfsFullfn[len("/boot/initramfs-"):]
            initramfsFilesFullfn = "/boot/initramfs-files-%s.tar.bz2" % (postfix)
            if not os.path.exists(initramfsFilesFullfn):
                die('File "%s" does not exist.' % (initramfsFilesFullfn))
            i += 1

        op = sys.argv[i]
        if op == "add-bin":
            if len(sys.argv) < i + 2:
                usage(True)
            path = sys.argv[i + 1]
            if not path.startswith("/"):
                die('The "path" parameter must be an absolute path.')
        elif op == "add-file":
            if len(sys.argv) < i + 3:
                usage(True)
            path = sys.argv[i + 1]
            targetPath = sys.argv[i + 2]
            if not targetPath.startswith("/"):
                die('The "target" parameter must be an absolute path.')
        else:
            usage(True)

        # find initramfs file if neccessary
        if initramfsFullfn is None:
            ret = []
            for kernelFile in sorted(os.listdir("/boot"), reverse=True):
                if kernelFile.startswith("kernel-"):
                    postfix = kernelFile[len("kernel-"):]
                    fullfn = "/boot/initramfs-%s" % (postfix)
                    fullfn2 = "/boot/initramfs-files-%s.tar.bz2" % (postfix)
                    if os.path.exists(fullfn):
                        ret.append((fullfn, fullfn2))
            if len(ret) == 0:
                die('No initramfs file found in "/boot".')
            elif len(ret) > 1:
                die('Multiple initramfs file found in "/boot".')
            else:
                initramfsFullfn, initramfsFilesFullfn = ret[0]
                if not os.path.exists(initramfsFilesFullfn):
                    die('File "%s" exists, but "%s" does not exist.' % (initramfsFullfn, initramfsFilesFullfn))

    with tempfile.TemporaryDirectory() as tmpDir:
        # decompress
        with tarfile.open(initramfsFilesFullfn, "r:bz2") as f:
            f.extractall(tmpDir)

        # do work
        if op == "add-bin":
            Util.copyWithSymlink("/", tmpDir, path[1:])
            for df in Util.libUsed(path):
                Util.copyWithSymlink("/", tmpDir, df[1:])
        elif op == "add-file":
            dstfile = os.path.join(tmpDir, targetPath[1:])
            dstdir = os.path.dirname(dstfile)
            if not os.path.exists(dstdir):
                os.makedirs(dstdir)
            subprocess.check_call(["cp", "-f", path, dstfile])
        else:
            assert False

        # prepare /boot
        mp = PhysicalDiskMounts.find_entry_by_mount_point("/boot")
        if mp is not None and "rw" not in mp.mnt_opt_list:
            subprocess.check_call(["mount", "/boot", "-o", "remount,rw"])

        # re-compress
        with TempChdir(tmpDir):
            cmdStr = "find . -print0 | cpio --null -H newc -o "
            cmdStr += "| xz --format=lzma "
            cmdStr += "> \"%s\" " % (initramfsFullfn)
            subprocess.check_call(cmdStr, shell=True, text=True)

            with tarfile.open(initramfsFilesFullfn, "w:bz2") as f:
                for fn in glob.glob("*"):
                    f.add(fn)
